# -*- coding: utf-8 -*-
"""
Author: chen xi (modified by Qwen)
Date: 2024/12/24 下午2:30
File: mass_tcp_client_random_string.py
功能：高并发 TCP 客户端，发送随机长度的加密二进制数据，接收时不尝试 UTF-8 解码
"""

import asyncio
import random
import string

from string_enc import iPN_StringEnc


def generate_random_string(min_len=1, max_len=64):
    """生成随机长度的随机字符串"""
    length = random.randint(min_len, max_len)
    chars = string.ascii_letters + string.digits + "!@#$%^&*"
    return ''.join(random.choices(chars, k=length))


async def connect_and_send_periodically(host, port, interval):
    try:
        reader, writer = await asyncio.open_connection(host, port)
        addr = writer.get_extra_info('peername')
        print(f"Connected to {addr}")

        async def send_data_task():
            while True:
                await asyncio.sleep(interval)
                try:
                    # ✅ 每次发送前生成新的随机字符串
                    random_str = generate_random_string()
                    print(f"Sending '{random_str}...' to {addr}")
                    _, encrypted_data = iPN_StringEnc("RSIOTGWakeupKey@", random_str, 256)

                    # 确保是 bytes 类型
                    # ✅ 发送加密后的二进制数据（message 已是 bytes 或 str）
                    if isinstance(encrypted_data, str):
                        # 如果 message 是字符串（不应如此），先 encode
                        encrypted_data = encrypted_data.encode('utf-8')

                    writer.write(encrypted_data)
                    print(f"Sent {len(encrypted_data)} bytes to {addr} (from '{random_str[:32]}...')")
                    await writer.drain()

                except Exception as e:
                    print(f"Send error to {addr}: {e}")
                    break  # 发送失败则退出

        # 启动发送任务
        asyncio.create_task(send_data_task())

        # 接收循环：只读原始 bytes，不 decode
        try:
            while True:
                data = await reader.read(1024)
                if not data:
                    break
                print(f"Received {len(data)} bytes from {addr}: {data.hex()}")
        except Exception as e:
            print(f"Error receiving from {addr}: {type(e).__name__}: {e}")
        finally:
            print(f"Closing connection to {addr}")
            if not writer.is_closing():
                writer.close()
                await writer.wait_closed()

    except Exception as e:
        print(f"Failed to connect to {host}:{port}: {e}")


async def create_connections(server_id, num_connections, host, port, interval):
    tasks = []

    for i in range(server_id * num_connections, (server_id + 1) * num_connections):
        task = asyncio.create_task(
            connect_and_send_periodically(host, port, interval)
        )
        tasks.append(task)

        # 控制连接创建速率，避免瞬时风暴
        if i % 1000 == 0 and i > 0:
            await asyncio.sleep(0.5)

    await asyncio.gather(*tasks)


if __name__ == '__main__':
    num_connections = 1      # 每个客户端模拟的连接数
    server_id = 0                # 区分不同机器的 ID
    host = '47.105.76.209'        # 目标服务器 IP
    port = 50000                 # 目标端口
    interval = 1               # 发送间隔（秒）

    # 运行客户端
    asyncio.run(create_connections(server_id, num_connections, host, port, interval))